Hugging Faceのモデル学習で、各レイヤ毎に別のLearning Rateで学習する方法

Hugging Faceのモデル学習で、各レイヤ毎に別のLearning Rateで学習する方法

#Python
#機械学習
#Deep Learning
#Hugging Face
#自然言語処理
nokomoro3

nokomoro3

facebook logohatena logotwitter logo
Clock Icon2022.11.07

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

こんちには。

データアナリティクス事業本部 機械学習チームの中村です。

Hugging Faceのライブラリの使い方紹介記事第４弾です。

今回は、各レイヤ毎に別のLearning Rateで学習するためには、どのようにすればよいかを理解していきます。

前回のSchedulerと似た部分の内容になり、コードも一部流用しているため、併せて前回記事も参照ください。

層毎に別のLearning Rateで学習するとは

代表的な例としては、事前学習済みのBERT部分と、分類ヘッドで学習ヘッドを変更する例があります。

  • 事前学習済みのBERT部分：Lerning Rateを小さく設定し、学習時に少しだけ更新する。
  • 分類ヘッド：Lerning Rateを大きく設定し、学習時により大きく更新する。

今回は、この例にそって実装する方法を見ていきます。

実行環境

今回はGoogle Colaboratory環境で実行しました。

ハードウェアなどの情報は以下の通りです。

  • GPU: Tesla P100 (GPUメモリ16GB搭載)
  • CUDA: 11.1
  • メモリ: 13GB

主なライブラリのバージョンは以下となります。

  • transformers: 4.24.0
  • datasets: 2.6.1

インストール

transformersとdatasetsをインストールします。

!pip install transformers datasets

また事前学習モデルの依存モジュールをインストールします。

!pip install fugashi
!pip install ipadic
!pip install sentencepiece

ベースとするコード

今回のベースとするコードは、前回を踏襲した以下のコードとします。

49-55行目の部分がSchedulerに関する部分でした。今回はLearning Rateをすべてのstepで一定とするconstantをSchedulerにしています。

from datasets import load_dataset
from transformers import AutoTokenizer
from transformers import AutoModelForSequenceClassification
from transformers import TrainingArguments
from transformers import Trainer
from transformers import AdamW, get_constant_schedule
from sklearn.metrics import accuracy_score, f1_score
import torch

# データセットのロード
dataset = load_dataset("tyqiangz/multilingual-sentiments", "japanese")

# # 実験のためデータセットを縮小したい場合はコチラを有効化
# from datasets import DatasetDict
# dataset = DatasetDict({
#     "train": dataset['train'].select(range(100)),
#     "validation": dataset['validation'].select(range(100)),
#     "test": dataset['test'].select(range(100)),
# })

# トークナイザのロード
model_ckpt = "cl-tohoku/bert-base-japanese-whole-word-masking"
tokenizer = AutoTokenizer.from_pretrained(model_ckpt)

# トークナイズ処理
def tokenize(batch):
    return tokenizer(batch["text"], padding=True, truncation=True)
dataset_encoded = dataset.map(tokenize, batched=True, batch_size=None)

# 事前学習モデルのロード
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
num_labels = 3
model = (AutoModelForSequenceClassification
    .from_pretrained(model_ckpt, num_labels=num_labels)
    .to(device))

# メトリクスの定義
def compute_metrics(pred):
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)
    f1 = f1_score(labels, preds, average="weighted")
    acc = accuracy_score(labels, preds)
    return {"accuracy": acc, "f1": f1}

# 学習パラメータの設定
batch_size = 16
model_name = "sample-text-classification-bert"

# modelから学習すべきパラメータを抽出
params = filter(lambda x: x.requires_grad, model.parameters())

# 今回はoptimizerにAdamWを使用
optimizer = AdamW(params, lr=2e-5)

scheduler = get_constant_schedule(optimizer)

training_args = TrainingArguments(
    output_dir=model_name,
    num_train_epochs=10,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    weight_decay=0.01,
    evaluation_strategy="epoch",
    logging_strategy="steps",
    disable_tqdm=False,
    logging_steps=1,
    push_to_hub=False,
    log_level="error",
)

# Trainerの定義
trainer = Trainer(
    model=model,
    args=training_args,
    compute_metrics=compute_metrics,
    train_dataset=dataset_encoded["train"],
    eval_dataset=dataset_encoded["validation"],
    tokenizer=tokenizer,
    optimizers=[optimizer, scheduler],
)

# トレーニング実行
trainer.train()

この内容についての解説は前回の記事を参照ください。

各レイヤでLearning Rateを変更する方法

ポイントは、ベースとするコードにもある以下の部分です。

# modelから学習すべきパラメータを抽出
params = filter(lambda x: x.requires_grad, model.parameters())

ここで、今はmodel全体のパラメータを全て抽出していますが、これをレイヤごとに分けて取得できれば、異なるLearning Rateを適用することができます。

各パラメータの確認

model.parameters()では、各パラメータの名前が分かりませんが、代わりにmodel.named_parameters()を使用することでパラメータ名を確認できます。

list(name for name, param in model.named_parameters())

出力は以下となります。（長いため一部省略しています）

['bert.embeddings.word_embeddings.weight',
 'bert.embeddings.position_embeddings.weight',
 'bert.embeddings.token_type_embeddings.weight',
 'bert.embeddings.LayerNorm.weight',
 'bert.embeddings.LayerNorm.bias',
 'bert.encoder.layer.0.attention.self.query.weight',
 'bert.encoder.layer.0.attention.self.query.bias',
 'bert.encoder.layer.0.attention.self.key.weight',
 'bert.encoder.layer.0.attention.self.key.bias',
 'bert.encoder.layer.0.attention.self.value.weight',
 'bert.encoder.layer.0.attention.self.value.bias',
 'bert.encoder.layer.0.attention.output.dense.weight',
 'bert.encoder.layer.0.attention.output.dense.bias',
 'bert.encoder.layer.0.attention.output.LayerNorm.weight',
 'bert.encoder.layer.0.attention.output.LayerNorm.bias',
 'bert.encoder.layer.0.intermediate.dense.weight',
 'bert.encoder.layer.0.intermediate.dense.bias',
 'bert.encoder.layer.0.output.dense.weight',
 'bert.encoder.layer.0.output.dense.bias',
 'bert.encoder.layer.0.output.LayerNorm.weight',
 'bert.encoder.layer.0.output.LayerNorm.bias',
 'bert.encoder.layer.1.attention.self.query.weight',
 'bert.encoder.layer.1.attention.self.query.bias',
 'bert.encoder.layer.1.attention.self.key.weight',
 'bert.encoder.layer.1.attention.self.key.bias',
 'bert.encoder.layer.1.attention.self.value.weight',
 'bert.encoder.layer.1.attention.self.value.bias',
 'bert.encoder.layer.1.attention.output.dense.weight',
 'bert.encoder.layer.1.attention.output.dense.bias',
 'bert.encoder.layer.1.attention.output.LayerNorm.weight',
 'bert.encoder.layer.1.attention.output.LayerNorm.bias',
 'bert.encoder.layer.1.intermediate.dense.weight',
 'bert.encoder.layer.1.intermediate.dense.bias',
 'bert.encoder.layer.1.output.dense.weight',
 'bert.encoder.layer.1.output.dense.bias',
 'bert.encoder.layer.1.output.LayerNorm.weight',
 'bert.encoder.layer.1.output.LayerNorm.bias',
 'bert.encoder.layer.2.attention.self.query.weight',
 'bert.encoder.layer.2.attention.self.query.bias',
 'bert.encoder.layer.2.attention.self.key.weight',
 'bert.encoder.layer.2.attention.self.key.bias',
 'bert.encoder.layer.2.attention.self.value.weight',
 'bert.encoder.layer.2.attention.self.value.bias',
 'bert.encoder.layer.2.attention.output.dense.weight',
 'bert.encoder.layer.2.attention.output.dense.bias',
 'bert.encoder.layer.2.attention.output.LayerNorm.weight',
 'bert.encoder.layer.2.attention.output.LayerNorm.bias',
 'bert.encoder.layer.2.intermediate.dense.weight',
 'bert.encoder.layer.2.intermediate.dense.bias',
 'bert.encoder.layer.2.output.dense.weight',
 'bert.encoder.layer.2.output.dense.bias',
 'bert.encoder.layer.2.output.LayerNorm.weight',
 'bert.encoder.layer.2.output.LayerNorm.bias',
 
 ...中略...

 'bert.encoder.layer.11.attention.self.query.weight',
 'bert.encoder.layer.11.attention.self.query.bias',
 'bert.encoder.layer.11.attention.self.key.weight',
 'bert.encoder.layer.11.attention.self.key.bias',
 'bert.encoder.layer.11.attention.self.value.weight',
 'bert.encoder.layer.11.attention.self.value.bias',
 'bert.encoder.layer.11.attention.output.dense.weight',
 'bert.encoder.layer.11.attention.output.dense.bias',
 'bert.encoder.layer.11.attention.output.LayerNorm.weight',
 'bert.encoder.layer.11.attention.output.LayerNorm.bias',
 'bert.encoder.layer.11.intermediate.dense.weight',
 'bert.encoder.layer.11.intermediate.dense.bias',
 'bert.encoder.layer.11.output.dense.weight',
 'bert.encoder.layer.11.output.dense.bias',
 'bert.encoder.layer.11.output.LayerNorm.weight',
 'bert.encoder.layer.11.output.LayerNorm.bias',
 'bert.pooler.dense.weight',
 'bert.pooler.dense.bias',
 'classifier.weight',
 'classifier.bias']

簡単に説明すると、以下の部分が先頭の各種埋め込み層とLayerNormになります。

['bert.embeddings.word_embeddings.weight',
 'bert.embeddings.position_embeddings.weight',
 'bert.embeddings.token_type_embeddings.weight',
 'bert.embeddings.LayerNorm.weight',
 'bert.embeddings.LayerNorm.bias',

以下はEncoderの１層分です。baseではこれが12層あるため、layer.0～layer.11まであります。

 'bert.encoder.layer.0.attention.self.query.weight',
 'bert.encoder.layer.0.attention.self.query.bias',
 'bert.encoder.layer.0.attention.self.key.weight',
 'bert.encoder.layer.0.attention.self.key.bias',
 'bert.encoder.layer.0.attention.self.value.weight',
 'bert.encoder.layer.0.attention.self.value.bias',
 'bert.encoder.layer.0.attention.output.dense.weight',
 'bert.encoder.layer.0.attention.output.dense.bias',
 'bert.encoder.layer.0.attention.output.LayerNorm.weight',
 'bert.encoder.layer.0.attention.output.LayerNorm.bias',
 'bert.encoder.layer.0.intermediate.dense.weight',
 'bert.encoder.layer.0.intermediate.dense.bias',
 'bert.encoder.layer.0.output.dense.weight',
 'bert.encoder.layer.0.output.dense.bias',
 'bert.encoder.layer.0.output.LayerNorm.weight',
 'bert.encoder.layer.0.output.LayerNorm.bias',

以下はEncoderの最終層の先頭(CLSトークン相当部分)に線形層を通すpoolerの部分です。

 'bert.pooler.dense.weight',
 'bert.pooler.dense.bias',

最後に分類のためのヘッド(線形層)があります。

 'classifier.weight',
 'classifier.bias']

構造の詳細は、BERTの元論文やTransformerの解説をご覧ください。

設定方法

事前学習モデルの部分には"bert"という文字列、ヘッドには"classifier"という文字列で見分けて設定します。

body_params \
    = [p for n,p in model.named_parameters() if n.startswith("bert") and p.requires_grad]
head_params \
    = [p for n,p in model.named_parameters() if n.startswith("classifier") and p.requires_grad]

# 別々にLearning Rateを設定するには以下のようにする
optimizer = AdamW([
    {"params": body_params, 'lr': 0.0}, 
    {"params": head_params, 'lr': 2e-5}
])

scheduler = get_constant_schedule(optimizer)

動作確認

動作確認をします。念のため、先頭と末尾のパラメータをバックアップして学習前後の変化を確認します。

from datasets import load_dataset
from transformers import AutoTokenizer
from transformers import AutoModelForSequenceClassification
from transformers import TrainingArguments
from transformers import Trainer
from transformers import AdamW, get_constant_schedule
from sklearn.metrics import accuracy_score, f1_score
import torch

# データセットのロード
dataset = load_dataset("tyqiangz/multilingual-sentiments", "japanese")

# # 実験のためデータセットを縮小したい場合はコチラを有効化
# from datasets import DatasetDict
# dataset = DatasetDict({
#     "train": dataset['train'].select(range(100)),
#     "validation": dataset['validation'].select(range(100)),
#     "test": dataset['test'].select(range(100)),
# })

# トークナイザのロード
model_ckpt = "cl-tohoku/bert-base-japanese-whole-word-masking"
tokenizer = AutoTokenizer.from_pretrained(model_ckpt)

# トークナイズ処理
def tokenize(batch):
    return tokenizer(batch["text"], padding=True, truncation=True)
dataset_encoded = dataset.map(tokenize, batched=True, batch_size=None)

# 事前学習モデルのロード
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
num_labels = 3
model = (AutoModelForSequenceClassification
    .from_pretrained(model_ckpt, num_labels=num_labels)
    .to(device))

backup1 = list(model.parameters())[0].to('cpu').detach().numpy().copy()
backup2 = list(model.parameters())[-1].to('cpu').detach().numpy().copy()

# メトリクスの定義
def compute_metrics(pred):
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)
    f1 = f1_score(labels, preds, average="weighted")
    acc = accuracy_score(labels, preds)
    return {"accuracy": acc, "f1": f1}

# 学習パラメータの設定
batch_size = 16
model_name = "sample-text-classification-bert"

# 事前学習モデルの部分には"bert"という文字列、ヘッドには"classifier"という文字列で見分けて設定
body_params = [p for n,p in model.named_parameters() if n.startswith("bert") and p.requires_grad]
head_params = [p for n,p in model.named_parameters() if n.startswith("classifier") and p.requires_grad]

# 別々にLearning Rateを設定するには以下のようにする
optimizer = AdamW([
    {"params": body_params, 'lr': 0}, 
    {"params": head_params, 'lr': 2e-5}
])

scheduler = get_constant_schedule(optimizer)

training_args = TrainingArguments(
    output_dir=model_name,
    num_train_epochs=10,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    weight_decay=0.01,
    evaluation_strategy="epoch",
    logging_strategy="steps",
    disable_tqdm=False,
    logging_steps=1,
    push_to_hub=False,
    log_level="error",
)

# Trainerの定義
trainer = Trainer(
    model=model,
    args=training_args,
    compute_metrics=compute_metrics,
    train_dataset=dataset_encoded["train"],
    eval_dataset=dataset_encoded["validation"],
    tokenizer=tokenizer,
    optimizers=[optimizer, scheduler],
)

# トレーニング実行
trainer.train()

以下のように、先頭の層は値が変わらないことを確認できました。

backup1 - list(model.parameters())[0].to('cpu').detach().numpy().copy()

array([[0., 0., 0., ..., 0., 0., 0.],
       [0., 0., 0., ..., 0., 0., 0.],
       [0., 0., 0., ..., 0., 0., 0.],
       ...,
       [0., 0., 0., ..., 0., 0., 0.],
       [0., 0., 0., ..., 0., 0., 0.],
       [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)

また以下のように、末尾の層はパラメータが更新されていることが分かります。

backup2 - list(model.parameters())[-1].to('cpu').detach().numpy().copy()

array([ 0.00139277,  0.00139004, -0.00139202], dtype=float32)

今回のように複数のLearning Rateを設定した場合、trainerのログには先頭のLearning Rateが記載されるようですので、ご注意ください。

train_log = [i for i in trainer.state.log_history if "loss" in i]
train_log

[{'loss': 0.9356, 'learning_rate': 0, 'epoch': 0.14, 'step': 1},
 {'loss': 0.9791, 'learning_rate': 0, 'epoch': 0.29, 'step': 2},
 {'loss': 0.9827, 'learning_rate': 0, 'epoch': 0.43, 'step': 3},
 {'loss': 0.9548, 'learning_rate': 0, 'epoch': 0.57, 'step': 4},
 {'loss': 0.8873, 'learning_rate': 0, 'epoch': 0.71, 'step': 5},
 {'loss': 0.9497, 'learning_rate': 0, 'epoch': 0.86, 'step': 6},
 {'loss': 0.844, 'learning_rate': 0, 'epoch': 1.0, 'step': 7},
 {'loss': 0.903, 'learning_rate': 0, 'epoch': 1.14, 'step': 8},
 {'loss': 0.9116, 'learning_rate': 0, 'epoch': 1.29, 'step': 9},
 {'loss': 0.8593, 'learning_rate': 0, 'epoch': 1.43, 'step': 10},
 {'loss': 0.9059, 'learning_rate': 0, 'epoch': 1.57, 'step': 11},
 {'loss': 0.8507, 'learning_rate': 0, 'epoch': 1.71, 'step': 12},
 {'loss': 0.8737, 'learning_rate': 0, 'epoch': 1.86, 'step': 13},
 ...

より進んだ応用例

より進んだ応用例としては以下も参考にされてください。weight_decayなども各層で変更している例となります。

まとめ

いかがでしたでしょうか？

実際にBERTなどのfine-tuningは、より大規模モデル(largeなど)になるにつれて学習が難しくなるタスクもあるようです。 そういった際に、Learning Rateの細かい調整が必要になるシーンがあるため記事にしてみました。

本記事がHugging Faceを使われる方の参考になれば幸いです。

Share this article

facebook logohatena logotwitter logo

関連記事

API Gateway + Lambda + Pollyでテキストを音声に変換して、Flutterから再生してみた

API Gateway + Lambda + Pollyでテキストを音声に変換して、Flutterから再生してみた

User avatar
sora
2024.12.18
簡単にサンプルデータを生成するPythonのライブラリ「faker」

簡単にサンプルデータを生成するPythonのライブラリ「faker」

User avatar
石川覚
2024.12.17
AWS CloudShell内のpsqlとPython(Psycopg)からAurora DSQLに接続してみた

AWS CloudShell内のpsqlとPython(Psycopg)からAurora DSQLに接続してみた

User avatar
quiver
2024.12.16
Gemini 2.0 Flash Live API （試験運用版）を使ってAIと日本語で会話してみた

Gemini 2.0 Flash Live API （試験運用版）を使ってAIと日本語で会話してみた

User avatar
石川覚
2024.12.14
Classmethod's white logo
Facebook's logo
X's logo
YouTube's logo

© Classmethod, Inc. All rights reserved.